import math
from collections import Counter

def distance(a, b):
    return math.sqrt(sum((a[i] - b[i]) ** 2 for i in range(len(a))))

def KNN_predict(train_data, train_labels, x, k):
    dists = []
    for xi, label in zip(train_data, train_labels):
        d = distance(x, xi)
        dists.append((d, label))


    dists.sort(key=lambda t: t[0])
    k_nearest = dists[:k]

    labels = [label for _, label in k_nearest]
    most_common = Counter(labels).most_common(1)[0][0]

    return most_common


def KNN(train_data, train_labels, test_data, k=3):
    predictions = []
    for x in test_data:
        pred = KNN_predict(train_data, train_labels, x, k)
        predictions.append(pred)
    return predictions


if __name__ == "__main__":
    train_data = [
        [1, 2], [2, 3], [3, 3],  
        [8, 8], [9, 8], [8, 9]    
    ]
    train_labels = ['A', 'A', 'A', 'B', 'B', 'B']

    test_data = [
        [2, 2],
        [9, 9]
    ]

    preds = KNN(train_data, train_labels, test_data, k=3)
    print("预测结果:", preds)
